#include "ExecutionEngine/cgra_sim_api.h"
#include <algorithm>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <string>
#include <regex>
#include <vector>

#define STR_EQ(X, Y) strcmp(X, Y) == 0

class CoutRedirect {
public:
    CoutRedirect(std::streambuf* new_buffer)
        : old_buf_(std::cout.rdbuf(new_buffer)) {}

    ~CoutRedirect() {
        std::cout.rdbuf(old_buf_);
    }

private:
    std::streambuf* old_buf_;
};

std::vector<uint8_t> extract_data(const std::string& input, int algo) {
    std::vector<uint8_t> result;
    // 匹配形如 [BCU X,OFIFO OUT ID :Y ]ZZZZ 的模式
    std::regex pattern(R"(\[BCU \d\,OFIFO OUT ID :\d \]([A-Za-z0-9]+))");
    
    auto begin = std::sregex_iterator(input.begin(), input.end(), pattern);
    auto end = std::sregex_iterator();

    for (auto it = begin; it != end; ++it) {
        std::smatch match = *it;
        std::string hex_str = match[1].str();
        std::transform(hex_str.begin(), hex_str.end(), hex_str.begin(), ::tolower);

        // 过滤全0的情况
        if (hex_str.find_first_not_of('0') == std::string::npos) {
            continue;
        }

        // 补足8位前导零
        if (hex_str.length() < 8) {
            hex_str = std::string(8 - hex_str.length(), '0') + hex_str;
        }

        // 将8字符十六进制字符串转为4字节
        std::vector<uint8_t> raw_vec;
        for (size_t i = 0; i < 8; i += 2) {
            uint8_t byte = static_cast<uint8_t>(
                std::stoul(hex_str.substr(i, 2), nullptr, 16));
            raw_vec.push_back(byte);
        }
        if(algo == CK_MD_MD5 || algo == CK_MD_AES){
          // reverse for md5
          std::reverse(raw_vec.begin(), raw_vec.end());
        }
        result.insert(result.end(), raw_vec.begin(), raw_vec.end());
    }

    return result;
}

std::string execute_command(const char* cmd) {
    std::string result;
    FILE* pipe = popen(cmd, "r");    // POSIX系统
    if (!pipe) {
        throw std::runtime_error("command failed");
    }

    char buffer[128];
    while (!feof(pipe)) {
        if (ferror(pipe)) {
            pclose(pipe);
            throw std::runtime_error("read failed");
        }
        size_t bytes_read = fread(buffer, 1, sizeof(buffer), pipe);
        if (bytes_read > 0) {
            result.append(buffer, bytes_read);
        }
    }

    // 关闭管道并获取命令的退出状态
    int exit_status = pclose(pipe);

    if (exit_status != 0) {
        throw std::runtime_error("exit pipe failed " + std::to_string(exit_status));
    }
    return result;
}

void print(uint8_t* x) {
  for(int i = 0; i<16; i++) {
    printf("%02hhx ", x[i]);
  }
  printf("\n");
}

extern "C"
void algo_run(const char* algo_name, uint8_t *plain_addr, uint32_t plain_len,
              uint8_t *cipher_addr, uint8_t* key_addr = nullptr, uint32_t key_len = 0) {
  uint8_t pad_data[4096];
  uint32_t padlen = 0;
  uint32_t size = 0;

  memset(pad_data, 0x0, 4096);
  // param of simulator_ctrl
  const char* input_file_name;

  int algo;
  if (STR_EQ(algo_name, "sm3")) {
    algo = CK_MD_SM3;
    hash_padding(CK_MD_SM3, plain_addr, plain_len, pad_data, &padlen);
    input_file_name = "sm3_plaintext_in.fifo";
  }
  else if (STR_EQ(algo_name, "sha1")){
    algo = CK_MD_SHA1;
    hash_padding(CK_MD_SHA1, plain_addr, plain_len, pad_data, &padlen);
    input_file_name = "fifo.fifo";
  }
  else if (STR_EQ(algo_name, "md5")) {
    algo = CK_MD_MD5;
    hash_padding(CK_MD_MD5, plain_addr, plain_len, pad_data, &padlen);
    input_file_name = "plaintext_in.fifo";
  }else if (STR_EQ(algo_name, "aes")) {
    algo = CK_MD_AES;
    pkcs7_padding(plain_addr, plain_len, pad_data, &padlen);
    input_file_name = "plaintext_infifo.fifo";
  }

  // LOAD FROM ENV VARIABLE
  const char *simlib_path_cstr = std::getenv("CGRA_SIM_LIB");
  if (simlib_path_cstr == nullptr) {
    std::cerr << "Please set env: CGRA_SIM_LIB." << std::endl;
    exit(1);
  }
  char algo_work_path[256];
  strcpy(algo_work_path, simlib_path_cstr);
  strcat(algo_work_path, "/sources/");
  strcat(algo_work_path, algo_name);
  
  char fifo_file_path[256];
  strcpy(fifo_file_path, algo_work_path);
  strcat(fifo_file_path, "/resource/");
  strcat(fifo_file_path, input_file_name);

  char convertor_lib_path[256];
  strcpy(convertor_lib_path, simlib_path_cstr);
  strcat(convertor_lib_path, "/lib/librpuconvert.so");

  char rpusim_lib_path[256];
  strcpy(rpusim_lib_path, simlib_path_cstr);
  strcat(rpusim_lib_path, "/lib/librpusim.so");

  // std::cout << "writing file: " << fifo_file_path << std::endl;

  // write fifo file
  FILE *fifo_file = fopen(fifo_file_path, "w");
  if (fifo_file) {
    int len = -1;
    if (algo_name == "aes"){
        len = padlen;
    }else{
        len = 4 * 4 * 4;
    }
    for (int i = 0; i < len; i += 4) {
      if (i % (4 * 4) == 0 && i != 0) {
        fprintf(fifo_file, "\n");
      }
      if (i % (4 * 4) == 0){
        fprintf(fifo_file, "0x00000000 ");
      }
      uint32_t num = (static_cast<uint8_t>(pad_data[i]) |
                          static_cast<uint8_t>(pad_data[i + 1]) << 8 |
                          static_cast<uint8_t>(pad_data[i + 2]) << 16 |
                          static_cast<uint8_t>(pad_data[i + 3]) << 24);
      fprintf(fifo_file, "%08X ", num);
    }
    fclose(fifo_file);
  }else{
    std::cerr << "can not open fifo file: " << fifo_file_path << std::endl;
    exit(1);
  }

  // write imd file
  if(algo == CK_MD_AES) {
    char imd_path[256];
    strcpy(imd_path, algo_work_path);
    strcat(imd_path, "/resource/imd.imd");
    FILE *imd_file = fopen(imd_path, "w");
    if(!imd_file) {
      std::cerr << "can not open imd file: " << imd_path << std::endl;
      exit(1);
    }
    for (uint32_t i = 0; i < key_len; i += 4) {
        unsigned int value = ((unsigned int)key_addr[i+3] << 24) |
                             ((unsigned int)key_addr[i+2] << 16) |
                             ((unsigned int)key_addr[i+1] << 8) |
                             (unsigned int)key_addr[i];
        fprintf(imd_file, "0x%08X 0x%08X\n", i, value);
    }
    fclose(imd_file);
  }

  std::string log;
  {

    // simulate
#ifndef EXE_DIR
    std::cerr << "simulator_ctrl failed" << std::endl;
    exit(1);
#else
      char command[512] = EXE_DIR;
      strcat(command, "/simulator ");
      strcat(command, algo_name);
      strcat(command, " ");
      strcat(command, rpusim_lib_path);
      strcat(command, " ");
      strcat(command, algo_work_path);
      
      // std::cout << "executing command: " << command << std::endl;
      log = execute_command(command);
#endif
  }

  auto res_vector = extract_data(log, algo);
  for(size_t i = 0; i < res_vector.size(); i++){
    cipher_addr[i] = res_vector[i];
  }
  // printf("cipher_addr:\n");
  // print(cipher_addr);
  return;
}

extern "C"
void _mlir_ciface_cgra_sha1_run(RankedMemrefType<uint8_t, 1>* input,
              RankedMemrefType<uint8_t, 1>* output, uint32_t plain_len) {
  algo_run("sha1", input->data, plain_len, output->data);
  return;
}

extern "C"
void _mlir_ciface_cgra_md5_run(RankedMemrefType<uint8_t, 1>* input,
              RankedMemrefType<uint8_t, 1>* output, uint32_t plain_len) {
  algo_run("md5", input->data, plain_len, output->data);
  return;
}

extern "C"
void _mlir_ciface_cgra_sm3_run(RankedMemrefType<uint8_t, 1>* input,
              RankedMemrefType<uint8_t, 1>* output, uint32_t plain_len) {
  algo_run("sm3", input->data, plain_len, output->data);
  return;
}

extern "C"
void _mlir_ciface_cgra_aes_run(RankedMemrefType<uint8_t, 1>* input,
              RankedMemrefType<uint8_t, 1>* output, uint32_t plain_len, RankedMemrefType<uint8_t, 1>* key) {
  algo_run("aes", input->data, plain_len, output->data, key->data, key->sizes[0]);
  return;
}